{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/rom1504/clip-retrieval/.env/lib/python3.8/site-packages/clip/clip.py:23: UserWarning: PyTorch version 1.7.1 or higher is recommended\n",
      "  warnings.warn(\"PyTorch version 1.7.1 or higher is recommended\")\n"
     ]
    }
   ],
   "source": [
    "from clip_retrieval.clip_inference.reader import FilesReader, WebdatasetReader\n",
    "from clip_retrieval.clip_inference.runner import Sampler\n",
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n",
    "from clip_retrieval.clip_inference.load_clip import load_clip\n",
    "images = \"test_images\"\n",
    "tars = \"test_tars\"\n",
    "folder = images\n",
    "\n",
    "batch_size=2\n",
    "num_prepro_workers=2\n",
    "_, preprocess = load_clip()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "output_partition_id 0\n",
      "hi\n",
      "['test_images/123_456.jpg', 'test_images/389_535.jpg']\n",
      "torch.Size([2, 3, 224, 224])\n",
      "hi\n",
      "['test_images/524_316.jpg']\n",
      "torch.Size([1, 3, 224, 224])\n",
      "output_partition_id 1\n",
      "hi\n",
      "['test_images/208_495.jpg', 'test_images/416_264.jpg']\n",
      "torch.Size([2, 3, 224, 224])\n",
      "output_partition_id 2\n",
      "hi\n",
      "['test_images/321_421.jpg', 'test_images/456_123.jpg']\n",
      "torch.Size([2, 3, 224, 224])\n"
     ]
    }
   ],
   "source": [
    "output_partition_count = 3\n",
    "for output_partition_id in range(output_partition_count):\n",
    "    print(\"output_partition_id\", output_partition_id)\n",
    "    sampler = Sampler(output_partition_id=output_partition_id, output_partition_count=output_partition_count)\n",
    "    reader = FilesReader(sampler, preprocess, folder, batch_size, num_prepro_workers, enable_text=False, enable_image=True, enable_metadata=False)\n",
    "    for i in reader:\n",
    "        print(\"hi\")\n",
    "        print(i['image_filename'])\n",
    "        print(i['image_tensor'].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# turn this into a test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "output_partition_id 0\n",
      "hi\n",
      "['123_456', '208_495', '321_421', '389_535']\n",
      "torch.Size([4, 3, 224, 224])\n",
      "output_partition_id 1\n",
      "hi\n",
      "['416_264', '456_123', '524_316']\n",
      "torch.Size([3, 3, 224, 224])\n"
     ]
    }
   ],
   "source": [
    "output_partition_count = 2\n",
    "tars = [\"test_tars/image1.tar\", \"test_tars/image2.tar\"]\n",
    "for output_partition_id in range(output_partition_count):\n",
    "    print(\"output_partition_id\", output_partition_id)\n",
    "    sampler = Sampler(output_partition_id=output_partition_id, output_partition_count=output_partition_count)\n",
    "    reader = WebdatasetReader(sampler, preprocess, tars, batch_size, num_prepro_workers, enable_text=False, enable_image=True, enable_metadata=False)\n",
    "    for i in reader:\n",
    "        print(\"hi\")\n",
    "        print(i['image_filename'])\n",
    "        print(i['image_tensor'].shape)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "output_partition_id 0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "136it [00:23,  8.78it/s]/home/rom1504/clip-retrieval/.env/lib/python3.8/site-packages/webdataset/handlers.py:34: UserWarning: OSError('image file is truncated (65 bytes not processed)')\n",
      "  warnings.warn(repr(exn))\n",
      "290it [00:47,  6.10it/s]\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[1;32m/home/rom1504/clip-retrieval/tests/test_clip_inference/playground.ipynb Cell 6'\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      <a href='vscode-notebook-cell://ssh-remote%2B3080/home/rom1504/clip-retrieval/tests/test_clip_inference/playground.ipynb#ch0000036vscode-remote?line=8'>9</a>\u001b[0m sampler \u001b[39m=\u001b[39m Sampler(output_partition_id\u001b[39m=\u001b[39moutput_partition_id, output_partition_count\u001b[39m=\u001b[39moutput_partition_count)\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B3080/home/rom1504/clip-retrieval/tests/test_clip_inference/playground.ipynb#ch0000036vscode-remote?line=9'>10</a>\u001b[0m reader \u001b[39m=\u001b[39m WebdatasetReader(sampler, preprocess, tars, batch_size, num_prepro_workers, enable_text\u001b[39m=\u001b[39m\u001b[39mFalse\u001b[39;00m, enable_image\u001b[39m=\u001b[39m\u001b[39mTrue\u001b[39;00m, enable_metadata\u001b[39m=\u001b[39m\u001b[39mFalse\u001b[39;00m)\n\u001b[0;32m---> <a href='vscode-notebook-cell://ssh-remote%2B3080/home/rom1504/clip-retrieval/tests/test_clip_inference/playground.ipynb#ch0000036vscode-remote?line=10'>11</a>\u001b[0m \u001b[39mfor\u001b[39;00m i \u001b[39min\u001b[39;00m tqdm(reader):\n\u001b[1;32m     <a href='vscode-notebook-cell://ssh-remote%2B3080/home/rom1504/clip-retrieval/tests/test_clip_inference/playground.ipynb#ch0000036vscode-remote?line=11'>12</a>\u001b[0m     \u001b[39mpass\u001b[39;00m\n",
      "File \u001b[0;32m~/clip-retrieval/.env/lib/python3.8/site-packages/tqdm/std.py:1180\u001b[0m, in \u001b[0;36mtqdm.__iter__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/tqdm/std.py?line=1176'>1177</a>\u001b[0m time \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_time\n\u001b[1;32m   <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/tqdm/std.py?line=1178'>1179</a>\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m-> <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/tqdm/std.py?line=1179'>1180</a>\u001b[0m     \u001b[39mfor\u001b[39;00m obj \u001b[39min\u001b[39;00m iterable:\n\u001b[1;32m   <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/tqdm/std.py?line=1180'>1181</a>\u001b[0m         \u001b[39myield\u001b[39;00m obj\n\u001b[1;32m   <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/tqdm/std.py?line=1181'>1182</a>\u001b[0m         \u001b[39m# Update and possibly print the progressbar.\u001b[39;00m\n\u001b[1;32m   <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/tqdm/std.py?line=1182'>1183</a>\u001b[0m         \u001b[39m# Note: does not call self.update(1) for speed optimisation.\u001b[39;00m\n",
      "File \u001b[0;32m~/clip-retrieval/clip_retrieval/clip_inference/reader.py:197\u001b[0m, in \u001b[0;36mWebdatasetReader.__iter__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    <a href='file:///~/clip-retrieval/clip_retrieval/clip_inference/reader.py?line=195'>196</a>\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m__iter__\u001b[39m(\u001b[39mself\u001b[39m):\n\u001b[0;32m--> <a href='file:///~/clip-retrieval/clip_retrieval/clip_inference/reader.py?line=196'>197</a>\u001b[0m   \u001b[39mfor\u001b[39;00m batch \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdataloader:\n\u001b[1;32m    <a href='file:///~/clip-retrieval/clip_retrieval/clip_inference/reader.py?line=197'>198</a>\u001b[0m     \u001b[39myield\u001b[39;00m batch\n",
      "File \u001b[0;32m~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py:521\u001b[0m, in \u001b[0;36m_BaseDataLoaderIter.__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=518'>519</a>\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_sampler_iter \u001b[39mis\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m    <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=519'>520</a>\u001b[0m     \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_reset()\n\u001b[0;32m--> <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=520'>521</a>\u001b[0m data \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_next_data()\n\u001b[1;32m    <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=521'>522</a>\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_num_yielded \u001b[39m+\u001b[39m\u001b[39m=\u001b[39m \u001b[39m1\u001b[39m\n\u001b[1;32m    <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=522'>523</a>\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_dataset_kind \u001b[39m==\u001b[39m _DatasetKind\u001b[39m.\u001b[39mIterable \u001b[39mand\u001b[39;00m \\\n\u001b[1;32m    <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=523'>524</a>\u001b[0m         \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_IterableDataset_len_called \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m \u001b[39mand\u001b[39;00m \\\n\u001b[1;32m    <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=524'>525</a>\u001b[0m         \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_num_yielded \u001b[39m>\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_IterableDataset_len_called:\n",
      "File \u001b[0;32m~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py:1186\u001b[0m, in \u001b[0;36m_MultiProcessingDataLoaderIter._next_data\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=1182'>1183</a>\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_process_data(data)\n\u001b[1;32m   <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=1184'>1185</a>\u001b[0m \u001b[39massert\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_shutdown \u001b[39mand\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_tasks_outstanding \u001b[39m>\u001b[39m \u001b[39m0\u001b[39m\n\u001b[0;32m-> <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=1185'>1186</a>\u001b[0m idx, data \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_get_data()\n\u001b[1;32m   <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=1186'>1187</a>\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_tasks_outstanding \u001b[39m-\u001b[39m\u001b[39m=\u001b[39m \u001b[39m1\u001b[39m\n\u001b[1;32m   <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=1187'>1188</a>\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_dataset_kind \u001b[39m==\u001b[39m _DatasetKind\u001b[39m.\u001b[39mIterable:\n\u001b[1;32m   <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=1188'>1189</a>\u001b[0m     \u001b[39m# Check for _IterableDatasetStopIteration\u001b[39;00m\n",
      "File \u001b[0;32m~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py:1152\u001b[0m, in \u001b[0;36m_MultiProcessingDataLoaderIter._get_data\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=1147'>1148</a>\u001b[0m     \u001b[39m# In this case, `self._data_queue` is a `queue.Queue`,. But we don't\u001b[39;00m\n\u001b[1;32m   <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=1148'>1149</a>\u001b[0m     \u001b[39m# need to call `.task_done()` because we don't use `.join()`.\u001b[39;00m\n\u001b[1;32m   <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=1149'>1150</a>\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m   <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=1150'>1151</a>\u001b[0m     \u001b[39mwhile\u001b[39;00m \u001b[39mTrue\u001b[39;00m:\n\u001b[0;32m-> <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=1151'>1152</a>\u001b[0m         success, data \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_try_get_data()\n\u001b[1;32m   <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=1152'>1153</a>\u001b[0m         \u001b[39mif\u001b[39;00m success:\n\u001b[1;32m   <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=1153'>1154</a>\u001b[0m             \u001b[39mreturn\u001b[39;00m data\n",
      "File \u001b[0;32m~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py:990\u001b[0m, in \u001b[0;36m_MultiProcessingDataLoaderIter._try_get_data\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=976'>977</a>\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m_try_get_data\u001b[39m(\u001b[39mself\u001b[39m, timeout\u001b[39m=\u001b[39m_utils\u001b[39m.\u001b[39mMP_STATUS_CHECK_INTERVAL):\n\u001b[1;32m    <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=977'>978</a>\u001b[0m     \u001b[39m# Tries to fetch data from `self._data_queue` once for a given timeout.\u001b[39;00m\n\u001b[1;32m    <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=978'>979</a>\u001b[0m     \u001b[39m# This can also be used as inner loop of fetching without timeout, with\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=986'>987</a>\u001b[0m     \u001b[39m# Returns a 2-tuple:\u001b[39;00m\n\u001b[1;32m    <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=987'>988</a>\u001b[0m     \u001b[39m#   (bool: whether successfully get data, any: data if successful else None)\u001b[39;00m\n\u001b[1;32m    <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=988'>989</a>\u001b[0m     \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m--> <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=989'>990</a>\u001b[0m         data \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_data_queue\u001b[39m.\u001b[39;49mget(timeout\u001b[39m=\u001b[39;49mtimeout)\n\u001b[1;32m    <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=990'>991</a>\u001b[0m         \u001b[39mreturn\u001b[39;00m (\u001b[39mTrue\u001b[39;00m, data)\n\u001b[1;32m    <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=991'>992</a>\u001b[0m     \u001b[39mexcept\u001b[39;00m \u001b[39mException\u001b[39;00m \u001b[39mas\u001b[39;00m e:\n\u001b[1;32m    <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=992'>993</a>\u001b[0m         \u001b[39m# At timeout and error, we manually check whether any worker has\u001b[39;00m\n\u001b[1;32m    <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=993'>994</a>\u001b[0m         \u001b[39m# failed. Note that this is the only mechanism for Windows to detect\u001b[39;00m\n\u001b[1;32m    <a href='file:///~/clip-retrieval/.env/lib/python3.8/site-packages/torch/utils/data/dataloader.py?line=994'>995</a>\u001b[0m         \u001b[39m# worker failures.\u001b[39;00m\n",
      "File \u001b[0;32m/usr/lib/python3.8/multiprocessing/queues.py:107\u001b[0m, in \u001b[0;36mQueue.get\u001b[0;34m(self, block, timeout)\u001b[0m\n\u001b[1;32m    <a href='file:///usr/lib/python3.8/multiprocessing/queues.py?line=104'>105</a>\u001b[0m \u001b[39mif\u001b[39;00m block:\n\u001b[1;32m    <a href='file:///usr/lib/python3.8/multiprocessing/queues.py?line=105'>106</a>\u001b[0m     timeout \u001b[39m=\u001b[39m deadline \u001b[39m-\u001b[39m time\u001b[39m.\u001b[39mmonotonic()\n\u001b[0;32m--> <a href='file:///usr/lib/python3.8/multiprocessing/queues.py?line=106'>107</a>\u001b[0m     \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_poll(timeout):\n\u001b[1;32m    <a href='file:///usr/lib/python3.8/multiprocessing/queues.py?line=107'>108</a>\u001b[0m         \u001b[39mraise\u001b[39;00m Empty\n\u001b[1;32m    <a href='file:///usr/lib/python3.8/multiprocessing/queues.py?line=108'>109</a>\u001b[0m \u001b[39melif\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_poll():\n",
      "File \u001b[0;32m/usr/lib/python3.8/multiprocessing/connection.py:257\u001b[0m, in \u001b[0;36m_ConnectionBase.poll\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    <a href='file:///usr/lib/python3.8/multiprocessing/connection.py?line=254'>255</a>\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_check_closed()\n\u001b[1;32m    <a href='file:///usr/lib/python3.8/multiprocessing/connection.py?line=255'>256</a>\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_check_readable()\n\u001b[0;32m--> <a href='file:///usr/lib/python3.8/multiprocessing/connection.py?line=256'>257</a>\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_poll(timeout)\n",
      "File \u001b[0;32m/usr/lib/python3.8/multiprocessing/connection.py:424\u001b[0m, in \u001b[0;36mConnection._poll\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    <a href='file:///usr/lib/python3.8/multiprocessing/connection.py?line=422'>423</a>\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m_poll\u001b[39m(\u001b[39mself\u001b[39m, timeout):\n\u001b[0;32m--> <a href='file:///usr/lib/python3.8/multiprocessing/connection.py?line=423'>424</a>\u001b[0m     r \u001b[39m=\u001b[39m wait([\u001b[39mself\u001b[39;49m], timeout)\n\u001b[1;32m    <a href='file:///usr/lib/python3.8/multiprocessing/connection.py?line=424'>425</a>\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mbool\u001b[39m(r)\n",
      "File \u001b[0;32m/usr/lib/python3.8/multiprocessing/connection.py:931\u001b[0m, in \u001b[0;36mwait\u001b[0;34m(object_list, timeout)\u001b[0m\n\u001b[1;32m    <a href='file:///usr/lib/python3.8/multiprocessing/connection.py?line=927'>928</a>\u001b[0m     deadline \u001b[39m=\u001b[39m time\u001b[39m.\u001b[39mmonotonic() \u001b[39m+\u001b[39m timeout\n\u001b[1;32m    <a href='file:///usr/lib/python3.8/multiprocessing/connection.py?line=929'>930</a>\u001b[0m \u001b[39mwhile\u001b[39;00m \u001b[39mTrue\u001b[39;00m:\n\u001b[0;32m--> <a href='file:///usr/lib/python3.8/multiprocessing/connection.py?line=930'>931</a>\u001b[0m     ready \u001b[39m=\u001b[39m selector\u001b[39m.\u001b[39;49mselect(timeout)\n\u001b[1;32m    <a href='file:///usr/lib/python3.8/multiprocessing/connection.py?line=931'>932</a>\u001b[0m     \u001b[39mif\u001b[39;00m ready:\n\u001b[1;32m    <a href='file:///usr/lib/python3.8/multiprocessing/connection.py?line=932'>933</a>\u001b[0m         \u001b[39mreturn\u001b[39;00m [key\u001b[39m.\u001b[39mfileobj \u001b[39mfor\u001b[39;00m (key, events) \u001b[39min\u001b[39;00m ready]\n",
      "File \u001b[0;32m/usr/lib/python3.8/selectors.py:415\u001b[0m, in \u001b[0;36m_PollLikeSelector.select\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    <a href='file:///usr/lib/python3.8/selectors.py?line=412'>413</a>\u001b[0m ready \u001b[39m=\u001b[39m []\n\u001b[1;32m    <a href='file:///usr/lib/python3.8/selectors.py?line=413'>414</a>\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m--> <a href='file:///usr/lib/python3.8/selectors.py?line=414'>415</a>\u001b[0m     fd_event_list \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_selector\u001b[39m.\u001b[39;49mpoll(timeout)\n\u001b[1;32m    <a href='file:///usr/lib/python3.8/selectors.py?line=415'>416</a>\u001b[0m \u001b[39mexcept\u001b[39;00m \u001b[39mInterruptedError\u001b[39;00m:\n\u001b[1;32m    <a href='file:///usr/lib/python3.8/selectors.py?line=416'>417</a>\u001b[0m     \u001b[39mreturn\u001b[39;00m ready\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "output_partition_count = 2\n",
    "batch_size=256\n",
    "num_prepro_workers=8\n",
    "from braceexpand import braceexpand\n",
    "from tqdm import tqdm\n",
    "tars = [i for i in braceexpand(\"pipe:aws s3 cp s3://laion-us-east-1/laion-data/laion2B-data/{000000..231348}.tar -\")]\n",
    "for output_partition_id in range(output_partition_count):\n",
    "    print(\"output_partition_id\", output_partition_id)\n",
    "    sampler = Sampler(output_partition_id=output_partition_id, output_partition_count=output_partition_count)\n",
    "    reader = WebdatasetReader(sampler, preprocess, tars, batch_size, num_prepro_workers, enable_text=False, enable_image=True, enable_metadata=False)\n",
    "    for i in tqdm(reader):\n",
    "        pass\n",
    "        #print(\"hi\")\n",
    "        #print(i['image_filename'])\n",
    "        #print(i['image_tensor'].shape)\n",
    "        #break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# turn this into a test as well"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# next save the tensors, and test the mapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "output_partition_id 0\n"
     ]
    }
   ],
   "source": [
    "import pickle\n",
    "output_partition_count = 1\n",
    "batch_size=2\n",
    "tars = [\"test_tars/image1.tar\", \"test_tars/image2.tar\"]\n",
    "for output_partition_id in range(output_partition_count):\n",
    "    print(\"output_partition_id\", output_partition_id)\n",
    "    sampler = Sampler(output_partition_id=output_partition_id, output_partition_count=output_partition_count)\n",
    "    reader = WebdatasetReader(sampler, preprocess, tars, batch_size, num_prepro_workers, enable_text=False, enable_image=True, enable_metadata=False)\n",
    "    for i, sample in enumerate(reader):\n",
    "      with open(\"test_tensors/{}.pkl\".format(i), \"wb\") as f:\n",
    "        pickle.dump(sample, f)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "from clip_retrieval.clip_inference.mapper import ClipMapper\n",
    "\n",
    "mapper = ClipMapper(enable_image=True, enable_text=False, enable_metadata=False, use_mclip=False, device=\"cpu\", clip_model=\"ViT-B/32\", use_jit=True, mclip_model=\"\")\n",
    "tensor_files = [i for i in os.listdir(\"test_tensors\")]\n",
    "for tensor_file in tensor_files:\n",
    "  with open(\"test_tensors/{}\".format(tensor_file), \"rb\") as f:\n",
    "    tensor = pickle.load(f)\n",
    "    sample = mapper(tensor)\n",
    "    assert sample[\"image_embs\"].shape[0] == tensor['image_tensor'].shape[0]\n",
    "    with open(\"test_embeddings/{}\".format(tensor_file), \"wb\") as f:\n",
    "      pickle.dump(sample, f)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# next save the predictions, and test the writter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from clip_retrieval.clip_inference.writer import NumpyWriter\n",
    "import numpy as np\n",
    "import pickle \n",
    "import tempfile\n",
    "with tempfile.TemporaryDirectory() as tmpdir:\n",
    "  writer = NumpyWriter(partition_id=0, output_folder=tmpdir, enable_text=False, enable_image=True, enable_metadata=False, write_batch_size=10)\n",
    "  embedding_files = [i for i in os.listdir(\"test_embeddings\")]\n",
    "  expected_shape = 0\n",
    "  for embedding_file in embedding_files:\n",
    "    with open(\"test_embeddings/{}\".format(embedding_file), \"rb\") as f:\n",
    "      embedding = pickle.load(f)\n",
    "      expected_shape += embedding[\"image_embs\"].shape[0]\n",
    "      writer(embedding)\n",
    "  writer.flush()\n",
    "\n",
    "  with open(tmpdir + \"/img_emb/img_emb_0.npy\", \"rb\") as f:\n",
    "    image_embs = np.load(f)\n",
    "    assert image_embs.shape[0] == expected_shape\n",
    "  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# next do a runner test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from clip_retrieval.clip_inference.runner import Runner, Sampler\n",
    "from clip_retrieval.clip_inference.reader import FilesReader, WebdatasetReader\n",
    "from clip_retrieval.clip_inference.mapper import ClipMapper\n",
    "from clip_retrieval.clip_inference.writer import NumpyWriter\n",
    "from clip_retrieval.clip_inference.load_clip import load_clip\n",
    "import os\n",
    "import numpy as np\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n",
    "import tempfile\n",
    "\n",
    "\n",
    "output_partition_count=2\n",
    "num_prepro_workers=8\n",
    "batch_size=2\n",
    "folder = \"test_images\"\n",
    "\n",
    "\n",
    "with tempfile.TemporaryDirectory() as tmpdir:\n",
    "  def reader_builder(sampler):\n",
    "    _, preprocess = load_clip()\n",
    "    return FilesReader(sampler, preprocess, folder, batch_size, num_prepro_workers, enable_text=False, enable_image=True, enable_metadata=False)\n",
    "\n",
    "  def mapper_builder():\n",
    "    return ClipMapper(enable_image=True, enable_text=False, enable_metadata=False, use_mclip=False, device=\"cpu\", clip_model=\"ViT-B/32\", use_jit=True, mclip_model=\"\")\n",
    "\n",
    "  def writer_builder(i):\n",
    "    return NumpyWriter(partition_id=i, output_folder=tmpdir, enable_text=False, enable_image=True, enable_metadata=False, write_batch_size=10)\n",
    "\n",
    "  runner = Runner(reader_builder=reader_builder, mapper_builder=mapper_builder, writer_builder=writer_builder, output_partition_count=output_partition_count)\n",
    "\n",
    "  runner(0)\n",
    "\n",
    "  with open(tmpdir + \"/img_emb/img_emb_0.npy\", \"rb\") as f:\n",
    "    image_embs = np.load(f)\n",
    "    assert image_embs.shape[0] == 4\n",
    "  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# next do a standalone distributor test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "from clip_retrieval.clip_inference.runner import Runner, Sampler\n",
    "from clip_retrieval.clip_inference.reader import FilesReader, WebdatasetReader\n",
    "from clip_retrieval.clip_inference.mapper import ClipMapper\n",
    "from clip_retrieval.clip_inference.writer import NumpyWriter\n",
    "from clip_retrieval.clip_inference.load_clip import load_clip\n",
    "from clip_retrieval.clip_inference.distributor import SequentialDistributor, PysparkDistributor\n",
    "import os\n",
    "import numpy as np\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n",
    "import tempfile\n",
    "\n",
    "\n",
    "output_partition_count=2\n",
    "num_prepro_workers=8\n",
    "batch_size=2\n",
    "folder = \"test_images\"\n",
    "distributor_kind = \"pyspark\"\n",
    "\n",
    "\n",
    "with tempfile.TemporaryDirectory() as tmpdir:\n",
    "  def reader_builder(sampler):\n",
    "    _, preprocess = load_clip()\n",
    "    return FilesReader(sampler, preprocess, folder, batch_size, num_prepro_workers, enable_text=False, enable_image=True, enable_metadata=False)\n",
    "\n",
    "  def mapper_builder():\n",
    "    return ClipMapper(enable_image=True, enable_text=False, enable_metadata=False, use_mclip=False, device=\"cpu\", clip_model=\"ViT-B/32\", use_jit=True, mclip_model=\"\")\n",
    "\n",
    "  def writer_builder(i):\n",
    "    return NumpyWriter(partition_id=i, output_folder=tmpdir, enable_text=False, enable_image=True, enable_metadata=False, write_batch_size=10)\n",
    "\n",
    "  runner = Runner(reader_builder=reader_builder, mapper_builder=mapper_builder, writer_builder=writer_builder, output_partition_count=output_partition_count)\n",
    "\n",
    "  if distributor_kind == \"sequential\":\n",
    "    distributor = SequentialDistributor(runner, output_partition_count)\n",
    "  elif distributor_kind == \"pyspark\":\n",
    "    from pyspark.sql import SparkSession  # pylint: disable=import-outside-toplevel\n",
    "    spark = (\n",
    "        SparkSession.builder.config(\"spark.driver.memory\", \"16G\")\n",
    "        .master(\"local[\" + str(2) + \"]\")\n",
    "        .appName(\"spark-stats\")\n",
    "        .getOrCreate()\n",
    "    )\n",
    "    distributor = PysparkDistributor(runner, output_partition_count)\n",
    "  distributor()\n",
    "\n",
    "  with open(tmpdir + \"/img_emb/img_emb_0.npy\", \"rb\") as f:\n",
    "    image_embs = np.load(f)\n",
    "    assert image_embs.shape[0] == 4\n",
    "  with open(tmpdir + \"/img_emb/img_emb_1.npy\", \"rb\") as f:\n",
    "    image_embs = np.load(f)\n",
    "    assert image_embs.shape[0] == 3\n",
    "  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# next to an end to end main test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n",
    "import tempfile\n",
    "from clip_retrieval.clip_inference.main import main\n",
    "\n",
    "\n",
    "num_prepro_workers=8\n",
    "batch_size=2\n",
    "input_dataset = \"test_images\"\n",
    "distributor_kind = \"pyspark\"\n",
    "\n",
    "\n",
    "with tempfile.TemporaryDirectory() as tmpdir:\n",
    "  from pyspark.sql import SparkSession  # pylint: disable=import-outside-toplevel\n",
    "  spark = (\n",
    "    SparkSession.builder.config(\"spark.driver.memory\", \"16G\")\n",
    "    .master(\"local[\" + str(2) + \"]\")\n",
    "    .appName(\"spark-stats\")\n",
    "    .getOrCreate()\n",
    "  )\n",
    "\n",
    "\n",
    "  main(input_dataset,\n",
    "    output_folder=tmpdir,\n",
    "    input_format=\"files\",\n",
    "    cache_path=None,\n",
    "    batch_size=256,\n",
    "    num_prepro_workers=8,\n",
    "    enable_text=False,\n",
    "    enable_image=True,\n",
    "    enable_metadata=False,\n",
    "    write_batch_size=4,\n",
    "    wds_image_key=\"jpg\",\n",
    "    wds_caption_key=\"txt\",\n",
    "    clip_model=\"ViT-B/32\",\n",
    "    mclip_model=\"sentence-transformers/clip-ViT-B-32-multilingual-v1\",\n",
    "    use_mclip=False,\n",
    "    use_jit=True,\n",
    "    distribution_strategy=\"pyspark\",\n",
    "    wds_number_file_per_input_file=10000,\n",
    "    output_partition_count=None,)\n",
    "\n",
    "\n",
    "  with open(tmpdir + \"/img_emb/img_emb_0.npy\", \"rb\") as f:\n",
    "    image_embs = np.load(f)\n",
    "    assert image_embs.shape[0] == 4\n",
    "  with open(tmpdir + \"/img_emb/img_emb_1.npy\", \"rb\") as f:\n",
    "    image_embs = np.load(f)\n",
    "    assert image_embs.shape[0] == 3\n",
    "\n",
    "  \n",
    "  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# next test with real data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# done"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test the logger writer and the logger reader here then in a test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/rom1504/clip-retrieval/.env/lib/python3.8/site-packages/clip/clip.py:23: UserWarning: PyTorch version 1.7.1 or higher is recommended\n",
      "  warnings.warn(\"PyTorch version 1.7.1 or higher is recommended\")\n"
     ]
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/rom1504/clip-retrieval/.env/lib/python3.8/site-packages/clip/clip.py:23: UserWarning: PyTorch version 1.7.1 or higher is recommended\n",
      "  warnings.warn(\"PyTorch version 1.7.1 or higher is recommended\")\n",
      "/home/rom1504/clip-retrieval/.env/lib/python3.8/site-packages/clip/clip.py:23: UserWarning: PyTorch version 1.7.1 or higher is recommended\n",
      "  warnings.warn(\"PyTorch version 1.7.1 or higher is recommended\")\n"
     ]
    }
   ],
   "source": [
    "from clip_retrieval.clip_inference.logger import LoggerWriter\n",
    "logger = LoggerWriter(partition_id=0, stats_folder=\"/tmp/my_stats\")\n",
    "logger.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " sample_per_sec 55086 ; sample_count 10358784 "
     ]
    }
   ],
   "source": [
    "for i in range(10000):\n",
    "  logger({  \"read_duration\": 0.5,\n",
    "            \"inference_duration\": 10,\n",
    "            \"write_duration\": 2,\n",
    "            \"total_duration\": 13,\n",
    "            \"sample_count\": 1024})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger.end()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from clip_retrieval.clip_inference.logger import LoggerReader\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/rom1504/clip-retrieval/.env/lib/python3.8/site-packages/clip/clip.py:23: UserWarning: PyTorch version 1.7.1 or higher is recommended\n",
      "  warnings.warn(\"PyTorch version 1.7.1 or higher is recommended\")\n",
      "wandb: Currently logged in as: rom1504 (use `wandb login --relogin` to force relogin)\n",
      "wandb: Tracking run with wandb version 0.12.10\n",
      "wandb: Syncing run silver-blaze-5\n",
      "wandb: ⭐️ View project at https://wandb.ai/rom1504/clip_retrieval\n",
      "wandb: 🚀 View run at https://wandb.ai/rom1504/clip_retrieval/runs/2wyv3p9i\n",
      "wandb: Run data is saved locally in /home/rom1504/clip-retrieval/tests/test_clip_inference/wandb/run-20220217_004542-2wyv3p9i\n",
      "wandb: Run `wandb offline` to turn off syncing.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " sample_per_sec 31 ; sample_count 1024  "
     ]
    }
   ],
   "source": [
    "reader = LoggerReader(stats_folder=\"/tmp/my_stats\", enable_wandb=True)\n",
    "reader.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>aa</th>\n",
       "      <th>b</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>23</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>23</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   aa   b\n",
       "0   1  23\n",
       "1   1  23"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "a = pd.read_parquet(\"/tmp/my_stats/stats.parquet\")\n",
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# then benchmark it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# locally done"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# then do the doc for pyspark with gpu setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# then do the real benchmark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# then done"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "8890164936ba431effa62f548d2e190a63033d8c51925a70e93a060bef4e9d5d"
  },
  "kernelspec": {
   "display_name": "Python 3.8.10 ('.env': venv)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
